import os
import json
import argparse
import logging
import random
import time
from datetime import datetime
from typing import Dict, Any, List, Set, Optional

# 导入自定义模块
from request_handler import RequestHandler
from parser import XiaohongshuParser
from note_parser import XiaohongshuNoteParser
from data_extractor import DataExtractor
from url_manager import URLManager
from config import XHSConfig

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('crawler.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('XiaohongshuCrawler')

# 主函数入口
def main():
    """
    主函数，处理命令行参数并执行爬虫任务
    """
    parser = argparse.ArgumentParser(description='小红书爬虫工具')
    
    # 核心功能参数组
    core_group = parser.add_argument_group('核心功能')
    core_group.add_argument('--keyword', '-k', type=str, help='搜索关键词')
    core_group.add_argument('--keywords', '-ks', type=str, help='搜索关键词列表，用逗号分隔')
    core_group.add_argument('--url', '-u', type=str, help='单个笔记URL')
    core_group.add_argument('--urls-file', '-uf', type=str, help='包含URL列表的文件路径')
    
    # 数据量控制参数
    count_group = parser.add_argument_group('数据量控制')
    count_group.add_argument('--target-count', '-tc', type=int, default=1000, 
                          help='目标爬取数量，默认1000条')
    count_group.add_argument('--max-pages', '-mp', type=int, default=100, 
                          help='最大搜索页数，默认100页')
    
    # 配置参数组
    config_group = parser.add_argument_group('配置选项')
    config_group.add_argument('--output-dir', '-o', type=str, default='output', 
                           help='输出目录路径')
    config_group.add_argument('--cookies-file', '-cf', type=str, default='cookies.json', 
                           help='cookies文件路径')
    config_group.add_argument('--use-proxy', '-p', action='store_true', 
                           help='是否使用代理')
    config_group.add_argument('--proxy-file', '-pf', type=str, default='proxies.txt', 
                           help='代理文件路径')
    config_group.add_argument('--crawl-comments', '-cc', action='store_true', 
                           help='是否爬取评论')
    
    # 排序选项
    parser.add_argument('--sort-by', '-sb', type=str, default='general', 
                      choices=['general', 'popular'], 
                      help='排序方式，默认general（综合），可选popular（热门）')
    
    args = parser.parse_args()
    
    # 验证参数
    if not any([args.keyword, args.keywords, args.url, args.urls_file]):
        parser.error('请至少提供一个关键词、关键词列表、单个URL或URL文件')
    
    # 创建配置
    config = XHSConfig()
    config.crawl_comments = args.crawl_comments
    
    # 创建爬虫实例
    crawler = XiaohongshuCrawler(
        cookies_file=args.cookies_file,
        output_dir=args.output_dir,
        config=config,
        use_proxy=args.use_proxy,
        proxy_file=args.proxy_file if args.use_proxy else None
    )
    
    try:
        # 根据参数执行不同的爬取任务
        if args.keyword:
            # 单个关键词爬取，支持指定目标数量
            print(f"\n========== 开始通过关键词爬取 ==========")
            print(f"关键词: {args.keyword}")
            print(f"目标数量: {args.target_count}")
            print(f"最大页数: {args.max_pages}")
            print(f"排序方式: {args.sort_by}")
            print("====================================\n")
            
            crawler.crawl_by_keyword_with_target_count(
                keyword=args.keyword,
                target_count=args.target_count,
                max_pages=args.max_pages,
                sort_by=args.sort_by
            )
            
        elif args.keywords:
            # 多个关键词爬取
            keywords_list = [k.strip() for k in args.keywords.split(',')]
            print(f"\n========== 开始通过多关键词爬取 ==========")
            print(f"关键词列表: {keywords_list}")
            print(f"目标总数量: {args.target_count}")
            print(f"最大页数: {args.max_pages}")
            print("========================================\n")
            
            # 为每个关键词分配目标数量
            per_keyword_count = args.target_count // len(keywords_list)
            remaining_count = args.target_count % len(keywords_list)
            
            for i, keyword in enumerate(keywords_list):
                # 最后一个关键词分配剩余的数量
                current_target = per_keyword_count + (remaining_count if i == len(keywords_list) - 1 else 0)
                print(f"\n处理关键词 {i+1}/{len(keywords_list)}: {keyword} (目标数量: {current_target})")
                crawler.crawl_by_keyword_with_target_count(
                    keyword=keyword,
                    target_count=current_target,
                    max_pages=args.max_pages,
                    sort_by=args.sort_by
                )
                
                # 不同关键词之间添加较长延迟
                if i < len(keywords_list) - 1:
                    delay = random.uniform(10, 20)
                    print(f"关键词'{keyword}'爬取完成，等待{delay:.1f}秒后处理下一个关键词")
                    time.sleep(delay)
                    
        elif args.url:
            # 单个URL爬取
            print(f"\n========== 开始爬取单个笔记 ==========")
            print(f"URL: {args.url}")
            print("====================================\n")
            
            result = crawler.crawl_note(args.url)
            if result.get('success'):
                print(f"笔记爬取成功: {args.url}")
            else:
                print(f"笔记爬取失败: {args.url}, 错误: {result.get('error')}")
                
        elif args.urls_file:
            # 从文件爬取多个URL
            print(f"\n========== 开始从文件爬取 ==========")
            print(f"文件路径: {args.urls_file}")
            print("==================================\n")
            
            if not os.path.exists(args.urls_file):
                print(f"错误: 文件不存在 - {args.urls_file}")
                return
                
            crawler.crawl_notes_from_file(args.urls_file)
            
        print("\n========== 爬取任务完成 ==========\n")
        
    except KeyboardInterrupt:
        print("\n爬虫任务被用户中断")
    except Exception as e:
        print(f"爬虫任务出错: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # 关闭爬虫资源
        crawler.close()
        print("爬虫资源已释放")
        
        # 生成执行报告
        print("\n========== 执行报告 ==========")
        print(f"爬取时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("爬虫任务已结束")
        print("==============================\n")
        
if __name__ == '__main__':
    main()


class XiaohongshuCrawler:
    """
    小红书爬虫主类，协调各个组件完成爬取任务
    """
    
    def __init__(self, cookies_file: str = 'cookies.json', output_dir: str = 'output', config: XHSConfig = None, use_proxy: bool = False, proxy_file: str = None):
        """
        初始化爬虫实例
        
        Args:
            cookies_file (str): cookies文件路径
            output_dir (str): 输出目录路径
            config (XHSConfig): 配置对象
            use_proxy (bool): 是否使用代理
            proxy_file (str): 代理文件路径
        """
        self.output_dir = output_dir
        self.cookies_file = cookies_file
        self.config = config or XHSConfig()
        
        # 创建输出目录
        self._create_output_dir()
        
        # 初始化组件
        self.request_handler = RequestHandler()
        
        # 加载cookies
        self._load_cookies()
        
        # 设置代理
        if use_proxy and proxy_file:
            self._load_proxies(proxy_file)
        
        # 初始化其他组件
        self.parser = XiaohongshuParser()
        self.note_parser = XiaohongshuNoteParser()
        self.data_extractor = DataExtractor()
        self.url_manager = URLManager()
        
        # 记录已爬取的笔记ID，避免重复爬取
        self.crawled_ids = set()
        self._load_crawled_ids()
    
    def _create_output_dir(self):
        """
        创建输出目录
        """
        if not os.path.exists(self.output_dir):
            try:
                os.makedirs(self.output_dir)
                logger.info(f"创建输出目录: {self.output_dir}")
            except Exception as e:
                logger.error(f"创建输出目录失败: {e}")
                raise
    
    def _load_cookies(self):
        """
        从文件加载cookies
        """
        if os.path.exists(self.cookies_file):
            try:
                success = self.request_handler.load_cookies(self.cookies_file)
                if success:
                    logger.info(f"成功加载cookies: {self.cookies_file}")
                else:
                    logger.warning(f"加载cookies失败: {self.cookies_file}")
            except Exception as e:
                logger.error(f"加载cookies时出错: {e}")
        else:
            logger.warning(f"cookies文件不存在: {self.cookies_file}")
    
    def crawl_note(self, note_url: str) -> Dict[str, Any]:
        """
        爬取单个笔记
        
        Args:
            note_url (str): 笔记URL
            
        Returns:
            dict: 爬取结果
        """
        try:
            # 提取笔记ID
            note_id = self.url_manager.extract_note_id(note_url)
            if not note_id:
                return {'success': False, 'error': f'无法从URL提取笔记ID: {note_url}'}
            
            # 检查是否已爬取
            if self._is_note_crawled(note_id):
                logger.info(f"笔记已爬取，跳过: {note_url}")
                return {'success': True, 'note_id': note_id, 'status': 'skipped'}
            
            # 获取笔记数据
            logger.info(f"正在获取笔记数据: {note_url}")
            note_data = self.request_handler.fetch_note_data(note_id)
            
            if not note_data:
                return {'success': False, 'error': f'获取笔记数据失败: {note_url}'}
            
            # 解析笔记数据
            logger.info(f"正在解析笔记数据: {note_id}")
            parsed_data = self.note_parser.parse(note_data)
            
            if not parsed_data:
                return {'success': False, 'error': f'解析笔记数据失败: {note_id}'}
            
            # 提取结构化数据
            structured_data = self.data_extractor.extract(parsed_data)
            
            # 保存爬取的数据
            self._save_crawled_data(note_id, structured_data)
            
            # 记录已爬取的ID
            self._save_crawled_id(note_id)
            
            logger.info(f"笔记爬取成功: {note_id}")
            return {'success': True, 'note_id': note_id, 'data': structured_data}
            
        except Exception as e:
            logger.error(f"爬取笔记时出错: {e}")
            return {'success': False, 'error': str(e)}
    
    def _save_crawled_data(self, note_id: str, data: Dict):
        """
        保存爬取的笔记数据
        
        Args:
            note_id (str): 笔记ID
            data (dict): 笔记数据
        """
        try:
            # 创建日期子目录
            date_str = datetime.now().strftime('%Y%m%d')
            date_dir = os.path.join(self.output_dir, date_str)
            if not os.path.exists(date_dir):
                os.makedirs(date_dir)
            
            # 保存数据到文件
            file_path = os.path.join(date_dir, f"{note_id}.json")
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            logger.info(f"数据保存成功: {file_path}")
            
        except Exception as e:
            logger.error(f"保存数据失败: {e}")
            raise
    
    def crawl_notes_from_file(self, urls_file: str):
        """
        从文件中批量爬取笔记
        
        Args:
            urls_file (str): 包含URL列表的文件路径
        """
        try:
            with open(urls_file, 'r', encoding='utf-8') as f:
                urls = [line.strip() for line in f if line.strip()]
            
            logger.info(f"从文件加载了 {len(urls)} 个URL")
            
            results = {
                'total': len(urls),
                'success': 0,
                'failed': 0,
                'skipped': 0
            }
            
            for i, url in enumerate(urls, 1):
                logger.info(f"正在爬取URL {i}/{len(urls)}: {url}")
                result = self.crawl_note(url)
                
                if result.get('success'):
                    if result.get('status') == 'skipped':
                        results['skipped'] += 1
                    else:
                        results['success'] += 1
                else:
                    results['failed'] += 1
                
                # 添加随机延迟
                if i < len(urls):
                    delay = random.uniform(2, 5)
                    logger.info(f"等待 {delay:.2f} 秒后继续")
                    time.sleep(delay)
            
            # 保存统计信息
            self._save_crawl_statistics(results)
            
        except Exception as e:
            logger.error(f"从文件爬取时出错: {e}")
            raise
    
    def _save_crawl_statistics(self, results: Dict):
        """
        保存爬取统计信息
        
        Args:
            results (dict): 爬取结果统计
        """
        try:
            stats_file = os.path.join(self.output_dir, f"crawl_stats_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
            with open(stats_file, 'w', encoding='utf-8') as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            
            logger.info(f"统计信息保存到: {stats_file}")
            
        except Exception as e:
            logger.error(f"保存统计信息失败: {e}")
    
    def _load_crawled_ids(self):
        """
        加载已爬取的笔记ID列表
        """
        try:
            # 检查所有日期目录
            if not os.path.exists(self.output_dir):
                return
            
            for date_dir in os.listdir(self.output_dir):
                date_path = os.path.join(self.output_dir, date_dir)
                if os.path.isdir(date_path):
                    # 读取所有JSON文件
                    for filename in os.listdir(date_path):
                        if filename.endswith('.json'):
                            note_id = filename[:-5]  # 移除.json后缀
                            self.crawled_ids.add(note_id)
            
            logger.info(f"已加载 {len(self.crawled_ids)} 个已爬取的笔记ID")
            
        except Exception as e:
            logger.error(f"加载已爬取ID时出错: {e}")
    
    def _save_crawled_id(self, note_id: str):
        """
        记录已爬取的笔记ID
        
        Args:
            note_id (str): 笔记ID
        """
        self.crawled_ids.add(note_id)
    
    def _is_note_crawled(self, note_id: str) -> bool:
        """
        检查笔记是否已爬取
        
        Args:
            note_id (str): 笔记ID
            
        Returns:
            bool: 是否已爬取
        """
        return note_id in self.crawled_ids
    
    def generate_search_url(self, keyword: str, page: int = 1) -> str:
        """
        生成搜索URL
        
        Args:
            keyword (str): 搜索关键词
            page (int): 页码
            
        Returns:
            str: 搜索URL
        """
        base_url = "https://www.xiaohongshu.com/search_result/"
        # 对关键词进行URL编码
        import urllib.parse
        encoded_keyword = urllib.parse.quote(keyword)
        return f"{base_url}?keyword={encoded_keyword}&page={page}"
    
    def search_notes(self, keyword: str, max_pages: int = None) -> List[str]:
        """
        搜索笔记并获取笔记ID列表
        
        Args:
            keyword (str): 搜索关键词
            max_pages (int): 最大爬取页数
            
        Returns:
            list: 笔记ID列表
        """
        if max_pages is None:
            max_pages = self.config.max_search_pages
        
        note_ids = []
        
        for page in range(1, max_pages + 1):
            logger.info(f"正在搜索第 {page}/{max_pages} 页")
            
            # 获取搜索结果页面
            search_url = self.generate_search_url(keyword, page)
            html_content = self.request_handler.get_page(search_url)
            
            if not html_content:
                logger.warning(f"无法获取第 {page} 页的搜索结果")
                continue
            
            # 解析搜索结果
            page_note_ids = self.parser.parse_search_results(html_content)
            
            if not page_note_ids:
                logger.info(f"第 {page} 页未找到笔记，可能已到达最后一页")
                break
            
            # 添加新的笔记ID
            new_note_ids = [note_id for note_id in page_note_ids if note_id not in note_ids]
            note_ids.extend(new_note_ids)
            
            logger.info(f"从第 {page} 页提取到 {len(new_note_ids)} 个新笔记ID")
            
            # 添加延迟
            delay = random.uniform(2, 5)
            logger.info(f"等待 {delay:.2f} 秒后继续")
            time.sleep(delay)
        
        logger.info(f"搜索完成，共找到 {len(note_ids)} 个笔记ID")
        return note_ids
    
    def crawl_by_keywords(self, keywords: List[str] = None, target_count: int = None):
        """
        通过关键词列表爬取笔记
        
        Args:
            keywords (list): 关键词列表
            target_count (int): 目标爬取数量
        """
        if keywords is None:
            keywords = self.config.default_keywords
        
        if target_count is None:
            target_count = self.config.default_target_count
        
        # 为每个关键词分配目标数量
        per_keyword_count = target_count // len(keywords)
        remaining_count = target_count % len(keywords)
        
        for i, keyword in enumerate(keywords):
            # 最后一个关键词分配剩余的数量
            current_target = per_keyword_count + (remaining_count if i == len(keywords) - 1 else 0)
            
            logger.info(f"开始爬取关键词 '{keyword}'，目标数量: {current_target}")
            
            # 搜索笔记
            note_ids = self.search_notes(keyword, max_pages=self.config.max_search_pages)
            
            # 过滤已爬取的笔记
            new_note_ids = [note_id for note_id in note_ids if note_id not in self.crawled_ids]
            
            # 如果新笔记数量超过目标，只取目标数量
            if len(new_note_ids) > current_target:
                new_note_ids = new_note_ids[:current_target]
            
            logger.info(f"找到 {len(new_note_ids)} 个新笔记需要爬取")
            
            # 爬取每个笔记
            for note_id in new_note_ids:
                note_url = f"https://www.xiaohongshu.com/explore/{note_id}"
                self.crawl_note(note_url)
                
                # 添加随机延迟
                delay = random.uniform(2, 5)
                time.sleep(delay)
            
            # 不同关键词之间添加较长延迟
            if i < len(keywords) - 1:
                delay = random.uniform(10, 20)
                logger.info(f"关键词 '{keyword}' 爬取完成，等待 {delay:.2f} 秒后处理下一个关键词")
                time.sleep(delay)
    
    def close(self):
        """
        关闭爬虫资源
        """
        try:
            # 保存cookies
            self.request_handler.save_cookies(self.cookies_file)
            
            # 关闭会话
            self.request_handler.close()
            
            logger.info("爬虫资源已释放")
            
        except Exception as e:
            logger.error(f"关闭爬虫资源时出错: {e}")
    
    # 增强版：直接使用request_handler的fetch_many_search_results方法获取大量数据
    def crawl_by_keyword_with_target_count(self, keyword: str, target_count: int = 1000, max_pages: int = 100, sort_by: str = 'general'):
        """
        通过关键词搜索并爬取指定数量的笔记
        直接调用底层fetch_many_search_results获取大量笔记ID
        
        Args:
            keyword (str): 搜索关键词
            target_count (int): 目标爬取数量
            max_pages (int): 最大搜索页数
            sort_by (str): 排序方式 ('general' 或 'popular')
        """
        logger.info(f"开始通过关键词'{keyword}'搜索并爬取 {target_count} 条笔记")
        
        # 第一步：使用增强的搜索方法获取大量笔记ID
        logger.info(f"正在搜索笔记ID，目标数量: {target_count}")
        note_ids = self.request_handler.fetch_many_search_results(
            keyword=keyword,
            max_pages=max_pages,
            max_note_ids=target_count,
            sort_by=sort_by
        )
        
        logger.info(f"搜索完成，共找到 {len(note_ids)} 个笔记ID")
        
        # 过滤掉已爬取的笔记ID
        new_note_ids = [note_id for note_id in note_ids if note_id not in self.crawled_ids]
        logger.info(f"过滤后有 {len(new_note_ids)} 个新笔记ID需要爬取")
        
        # 第二步：爬取每个笔记的详细信息
        stats = {
            'keyword': keyword,
            'total_ids': len(new_note_ids),
            'success': 0,
            'failed': 0,
            'failed_ids': []
        }
        
        for i, note_id in enumerate(new_note_ids, 1):
            logger.info(f"正在爬取笔记 {i}/{len(new_note_ids)}: {note_id}")
            note_url = f"https://www.xiaohongshu.com/explore/{note_id}"
            result = self.crawl_note(note_url)
            
            if result.get('success'):
                stats['success'] += 1
            else:
                stats['failed'] += 1
                stats['failed_ids'].append(note_id)
            
            # 添加随机延迟，避免被反爬
            if i < len(new_note_ids):
                delay = random.uniform(2, 5)
                logger.info(f"等待 {delay:.2f} 秒后继续")
                time.sleep(delay)
        
        # 保存统计信息
        stats_file = os.path.join(self.output_dir, f"keyword_stats_{keyword}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
        try:
            with open(stats_file, 'w', encoding='utf-8') as f:
                json.dump(stats, f, ensure_ascii=False, indent=2)
            logger.info(f"统计信息保存到: {stats_file}")
        except Exception as e:
            logger.error(f"保存统计信息失败: {e}")
        
        # 打印爬取结果摘要
        print(f"\n========== 关键词 '{keyword}' 爬取结果摘要 ==========")
        print(f"目标数量: {target_count}")
        print(f"搜索到的笔记ID总数: {len(note_ids)}")
        print(f"新笔记ID数量: {len(new_note_ids)}")
        print(f"成功爬取: {stats['success']}")
        print(f"爬取失败: {stats['failed']}")
        print(f"成功率: {(stats['success'] / len(new_note_ids) * 100) if new_note_ids else 0:.2f}%")
        print(f"统计信息文件: {stats_file}")
        print("================================================\n")
